"""
Mengji framework 通过 RIVER/ASche 两种部件组合而来，由 RIVER 实现数据流管理，ASche 实现任务执行时序管理。
Mengji framework 提供基本虚拟总线控制器（virtual_bus），单元处理器（unit_processor）和接线（wire）。
此文件实现了 RIVER 的管线结构和三大件基本的模版框架供功能扩展使用。

date: 2025/5/7
author: SiHeng Tang
file: task_element.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import asyncio
import gc
import logging
import time

import async_sche
from async_sche import ExitGenerationError
from mj_errors import MJRuntimeError

logger = logging.getLogger(__name__)

DATA_PACK_TIME_STAMP = "[%U:%w:%H:%M:%S]"


class DataPack(object):
    """
    RIVER 数据包，现阶段数据包只包含数据进程内交换和字符化打印功能。

    数据包包含了几个基本属性：
        - ``pack_time`` 微秒计数器，用于定位数据包时间位置
        - ``payload`` 有效荷载，数据的 Python 对象，在未来的跨进程传输中，此对象需要支持序列化和反序列化。
        - ``meta_data`` 元数据

    可以使用成员访存函数的对象：``pack_time`` ``payload`` 直接使用键名访问元数据
    """

    def __init__(self, **kwargs):
        """
        :param kwargs: 透传给 ``set_v`` 设置值
        """
        self._payload = None

        self._pack_time = time.time()
        self._pack_time_struct = time.gmtime(self._pack_time)
        self._meta_data = {}

        self.set_v(**kwargs)

    def __str__(self):
        _self_str = f"Data Pack "
        _self_str += f"{time.strftime(DATA_PACK_TIME_STAMP, self._pack_time_struct)}\n"
        _self_str += f"---- meta_data: {self._meta_data=} "
        _self_str += f"content: {self._payload=}"
        return _self_str

    def set_v(self, **kwargs):
        """
        变长参数键值对风格赋值函数
        :param kwargs: 使用 ``payload`` 键时赋值被应用到荷载，其他键赋值被应用到 ``meta_data``
        """
        for key, value in kwargs.items():
            if key == "payload":
                self._payload = value
            else:
                self._meta_data.setdefault(key, value)

    def get_d(self, *args) -> dict:
        """
        字典风格成员存取器
        :param args: 请求成员列表
        :return: 字典键值对，使用传参列表作为键，不存在的成员使用 None 替代
        """
        return_d = {}

        for key in args:
            if key == "payload":
                return_d.setdefault(key, self._payload)
            elif key == "pack_time":
                return_d.setdefault(key, self._pack_time)
            else:
                return_d.setdefault(key, self._meta_data.get(key))

        return return_d

    def get_s(self, *args) -> tuple:
        """
        列表风格成员存取器
        :param args: 请求成员列表
        :return: 按照请求成员顺序返回的列表，不存在的成员使用 None 替代
        """
        return tuple(self.get_d(*args).values())

    def get_a(self, k, default=None):
        """
        独立风格成员存取器
        :param k: 请求成员
        :param default: 默认返回值
        :return: 按照请求成员顺序返回的列表，不存在的成员使用 None 替代
        """
        return tuple(self.get_d(k).values())[0] if tuple(self.get_d(k).values())[0] else default

    # def encode_dp(self):
    #     pass
    #
    # def decode_dp(self):
    #     pass


class DataEndpoint(object):
    """
    数据端点通过 ``DataPipe`` 在两端点间双向传输数据，同时提供事件回调函数，实现响应式处理。
    数据端点重载 ``|`` 运算符提供数据交换，在数据交换前后各有一回调函数：
        - 交换前，before_trans
        - 交换后，after_trans
    回调函数声明::

        def handle_read(ep: DataEndpoint, method: str):
            # 回调内容
        ep.set_cb(before_out=func) # 注册函数
    """
    BEFORE_TRANS = "BEFORE_TRANS"
    AFTER_TRANS = "AFTER_TRANS"

    def __init__(self, **kwargs):
        """
        初始化端点
        :param kwargs: 透传给 ``set_cb`` 处理，设置回调函数
        """
        # 这里初始化空的数据包，解决应用启动时崩溃问题
        self._before_trans_cb = None
        self._after_trans_cb = None
        self._data_pack_out = DataPack()
        self._data_pack_in = DataPack()

        # 数据包更新监测
        self._updated_out_package = True
        self._updated_in_package = True

        self.set_cb(**kwargs)

    def set_cb(self, **kwargs):
        """
        :param kwargs: 回调函数键值对
        """
        self._before_trans_cb = kwargs["before_trans"] if "before_trans" in kwargs.keys() else self._before_trans_cb
        self._after_trans_cb = kwargs["after_trans"] if "after_trans" in kwargs.keys() else self._after_trans_cb

    def get_dp(self) -> DataPack:
        """
        :return: 返回输入的数据包
        """
        self._updated_in_package = False
        return self._data_pack_in

    def set_dp(self, new_dp: DataPack):
        """
        :param new_dp: 新数据包
        """
        self._data_pack_out = new_dp
        self._updated_out_package = True

    def detect_new_out_dp(self):
        """
        监测新数据包传出
        :return: 如果数据包更新，返回 True
        """
        return self._updated_out_package

    def detect_new_in_dp(self):
        """
        监测新数据包传入
        :return: 如果数据包更新，返回 True
        """
        return self._updated_in_package

    def __or__(self, other):
        """
        :param other: 另一侧的端点对象
        :return: (self.out, other.out)
        """
        # 这里引用同一类对象的保护方法没有问题
        if self._before_trans_cb: self._before_trans_cb(self, self.BEFORE_TRANS)
        if other._before_trans_cb: other._before_trans_cb(other, self.BEFORE_TRANS)

        self._data_pack_in = other._data_pack_out
        other._updated_out_package = False
        self._updated_in_package = True
        other._data_pack_in = self._data_pack_out
        self._updated_out_package = False
        other._updated_in_package = True

        if self._after_trans_cb: self._after_trans_cb(self, self.AFTER_TRANS)
        if other._after_trans_cb: other._after_trans_cb(other, self.AFTER_TRANS)

        return self._data_pack_out, other._data_pack_out


class DataPipe(object):
    """
    数据管线

    数据管线交换两个连接的端点中的数据，并且，要求每个端点的数据处理方法都是异步函数。
    """

    def __init__(self, **kwargs):
        """
        初始化数据管线并连接端点，使用 ``ep_a`` 和 ``ep_b`` 指定两个端点
        :param kwargs:
        """
        # 要求两个有效端点
        if (not "ep_a" in kwargs.keys()) or (not "ep_b" in kwargs.keys()):
            raise MJRuntimeError(MJRuntimeError.INTERNAL_ERROR, "Empty data endpoint in pipe side")
        if (type(kwargs["ep_a"]) != DataEndpoint) or (type(kwargs["ep_b"]) != DataEndpoint):
            raise MJRuntimeError(MJRuntimeError.INTERNAL_ERROR, "Invalid data endpoint type in pipe side")

        self._ep_b: DataEndpoint = kwargs["ep_b"]
        self._ep_a: DataEndpoint = kwargs["ep_a"]

    async def push(self):
        if self._ep_a and self._ep_b:
            self._ep_a | self._ep_b
        else:
            raise MJRuntimeError(MJRuntimeError.INTERNAL_ERROR, "Endpoint died when running")

    def disconnect(self):
        self._ep_a = None
        self._ep_b = None


class Wire(async_sche.BaseTaskGen, DataPipe):
    """
    接线，同时继承基本任务生成器的无限循环和数据管线的连接能力。
    """

    def __init__(self, **kwargs):
        """
        ``sche`` 调度器，``kwargs``
        :param sche:
        :param kwargs:
        """
        async_sche.BaseTaskGen.__init__(self, priority=async_sche.TaskWrapper.IDLE_PRIORITY + 20)
        DataPipe.__init__(self, **kwargs)

    def _on_gen(self):
        # 猜猜为什么新数据监测要在这里做而不是直接在端点传输时实现
        if not (self._ep_a.detect_new_out_dp() or self._ep_b.detect_new_out_dp()):
            raise ExitGenerationError()

    async def _func(self):
        await self.push()

    def _on_exit(self):
        self.disconnect()
        async_sche.BaseTaskGen._on_exit(self)


# #############
# 功能测试
# #############
# FIXME 过时的框架测试

class TestDataSrc(async_sche.BaseTaskGen):
    def __init__(self):
        super().__init__(priority=async_sche.TaskWrapper.IDLE_PRIORITY + 10)

        self.big_int = 0
        self.tele_gps_ep = DataEndpoint()
        self.tele_attitude_ep = DataEndpoint()
        self.tele_flight_mode_ep = DataEndpoint()

    async def _func(self):
        """
        收集遥测数据
        """
        gps_data = {
            "latitude": self.big_int,
            "longitude": self.big_int,
            "altitude": self.big_int
        }
        attitude_data = {
            "roll": self.big_int,
            "pitch": self.big_int,
            "yaw": self.big_int
        }
        flight_mode = 15

        self.tele_gps_ep.set_dp(DataPack(payload=gps_data))
        self.tele_attitude_ep.set_dp(DataPack(payload=attitude_data))
        self.tele_flight_mode_ep.set_dp(DataPack(payload=flight_mode))

        self.big_int += 1

        if self.big_int > 10:
            raise MJRuntimeError(MJRuntimeError.INTERNAL_ERROR, "test overflow exception")

    def _on_exit(self):
        self.tele_gps_ep = None
        self.tele_attitude_ep = None
        self.tele_flight_mode_ep = None


class PrintNode(async_sche.BaseTaskGen):
    def __init__(self):
        super().__init__(priority=async_sche.TaskWrapper.IDLE_PRIORITY + 10)

        self.recv_gps = DataEndpoint(after_trans=self.notify_by_ep)
        self.recv_attitude = DataEndpoint(after_trans=self.notify_by_ep)
        self.recv_flight_mode = DataEndpoint(after_trans=self.notify_by_ep)

    @staticmethod
    def notify_by_ep(ep, mtd):
        print(ep.get_dp())
        print(ep.get_dp().get_d("pack_time"))

    def _on_exit(self):
        self.recv_gps = None
        self.recv_attitude = None
        self.recv_flight_mode = None


async def test_part():
    sche = async_sche.MJScheduler()

    while True:
        plane = TestDataSrc()
        msg = PrintNode()
        w = Wire(ep_a=plane.tele_gps_ep, ep_b=msg.recv_gps)
        w1 = Wire(ep_a=plane.tele_attitude_ep, ep_b=msg.recv_attitude)
        w2 = Wire(ep_a=plane.tele_flight_mode_ep, ep_b=msg.recv_flight_mode)

        sche.reg_generator(plane)
        sche.reg_generator(msg)
        sche.reg_generator(w)
        sche.reg_generator(w1)
        sche.reg_generator(w2)

        try:
            await sche.main_loop()
        except MJRuntimeError:
            sche.finalize()

        plane = None
        msg = None
        w1 = None
        w2 = None
        w3 = None

        gc.collect()
        print(gc.garbage)


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    gc.set_debug(gc.DEBUG_LEAK)
    asyncio.run(test_part())
